#include <fstream>
#include "VolumeRadianceTransfer.h"

template<>
struct PositionTrait<Photon>
{
	static void getPosition(const Photon &v, Vector3 &p)
	{
		PositionTrait<Vector3>::getPosition(v.pos, p);
	}
};

template<>
struct HashTrait<Photon>
{
	static size_t hashCode(const Photon &key)
	{
		return HashTrait<Vector3>::hashCode(key.pos);
	}
};

template<>
struct EqualsTrait<Photon>
{
	static bool equals(const Photon &a, const Photon &b)
	{
		return EqualsTrait<Vector3>::equals(a.pos, b.pos);
	}
};

VolumeRadianceTransfer::VolumeRadianceTransfer()
{
	lmax = 4;
	psamples = 1;
	rsamples = 256;
	pbands = 2;
	rbands = 8;
	siga = 1.0f;
	sigs = 1.0f;
	minrad = -1.0f;
	maxrad = 1.0f;
	//pshader = Shader::fromFiles("PhotonVis.vp", "PhotonVis.gp", "PhotonVis.fp", 2);
}

void VolumeRadianceTransfer::tracePhotons(const Field<float> &field, int l, int m)
{
	// spherical coordinate
	// x-axis: phi = 0
	// y-axis: phi = pi/2
	// z-axis: theta = 0

	// get bounding box
	Box box(Vector3(0.0f, 0.0f, 0.0f), Vector3(field.nx, field.ny, field.nz) * field.dx);

	// initialize kd-tree
	std::cout << "building kd-tree ... ";
	kdtree.clear();

	float sig = siga + sigs;
	float pa = siga / sig;
	float ps = sigs / sig;
	float rad0 = 1.0f / psamples;

	// for each voxel
	for(int z = 0; z < field.nz; ++z) for(int y = 0; y < field.ny; ++y) for(int x = 0; x < field.nx; ++x)
	{
		// skip interior voxel
		if(x > 0 && x < field.nx - 1 && y > 0 && y < field.ny - 1 && z > 0 && z < field.nz - 1)
			continue;

		Vector3 pos = Vector3(x + 0.5f, y + 0.5f, z + 0.5f) * field.dx;

		// generate photons at boundary voxel
		//std::cout << "\rtracing voxel " << x << ", " << y << ", " << z << std::endl;
		for(int i = 0; i < psamples; ++i)
		{
			//std::cout << "\rtracing photon #" << i << std::endl;

			// sample random direction
			float theta, phi;
			SHFunction<float>::random(theta, phi);

			// emit photon
			float ylm = SHFunction<float>::realY(l, m, theta, phi);
			float cost = std::cos(theta);
			float sint = std::sin(theta);
			Photon p;
			p.pos = pos;
			p.dir = Vector3(cos(phi) * sint, sin(phi) * sint, cost) * -field.dx;
			p.rad = rad0 * ylm;

			// cache indirect lighting
			bool indirect = true;	// !!!! we just cache everything here !!!!

			// cumulative density
			float cd = 0.0f;

			while(box.contains(p.pos))
			{
				// integrate density
				float density = field.evaluate(p.pos.x, p.pos.y, p.pos.z);
				cd += (siga + sigs) * density * field.dx;

				float pr = Random::common().uniform();
				if(density >= 1e-3f && pr < 1.0f - exp(-cd))
				{
					pr = Random::common().uniform();
					if(pr < pa)
					{
						// absorption
						if(indirect)
							kdtree.insert(p);
						break;
					}
					else if(pr < pa + ps)
					{
						// scattering
						if(indirect)
							kdtree.insert(p);
						Random::common().sphere(p.dir.x, p.dir.y, p.dir.z);
						p.dir *= field.dx;
						indirect = true;
						cd = 0.0f;
					}
				}

				// integrate path
				p.pos += p.dir;
			}
		}
	}

	//kdtree.insert(photons);
	std::cout << "done!" << std::endl;
	std::cout << "balancing kd-tree ... ";
	kdtree.balance();
	std::cout << "done!" << std::endl;
	std::cout << "cached photons: " << kdtree.size() << std::endl;

	//minrad = 1e10f;
	//maxrad = -1e10f;
	//for(int i = 0; i < photons.size(); ++i)
	//{
	//	minrad = std::min(minrad, photons[i].rad);
	//	maxrad = std::max(maxrad, photons[i].rad);
	//}
	//if(minrad == maxrad)
	//{
	//	minrad -= 1e-3f;
	//	maxrad += 1e-3f;
	//}
}

bool intersectRayBox(const Ray &ray, const Vector3 &boxMin, const Vector3 &boxMax, float &t0, float &t1)
{
	Vector3 invDir = ray.invDirection();
	Vector3 tBot = invDir * (boxMin - ray.origin());
	Vector3 tTop = invDir * (boxMax - ray.origin());
	Vector3 tMin = G3D::min(tTop, tBot);
	Vector3 tMax = G3D::max(tTop, tBot);
	Vector2 temp = G3D::max(tMin.xx(), tMin.yz());
	t0 = G3D::max(temp.x, temp.y);
	temp = G3D::min(tMax.xx(), tMax.yz());
	t1 = G3D::min(temp.x, temp.y);
	if(t0 > t1 || t1 < 0.0f)
		return false;
	t0 = G3D::max(t0, 0.0f);
	return true;
}

float VolumeRadianceTransfer::estimateRadiance(const Field<float> &field, const Ray &ray, const Vector3 &boxMin, const Vector3 &boxMax)
{
	float t0, t1;
	bool hit = intersectRayBox(ray, boxMin, boxMax, t0, t1);
	if(!hit)
		return 0.0f;

	float dx = field.dx * 1.0f;
	float radius = field.dx * 2.0f;
	float volume = 4.0f / 3.0f * SHFunction<float>::PI * radius * radius * radius;
	float r2r2 = 0.5f / (radius * radius);
	float sig = siga + sigs;
	float sigdx = sig * dx;
	float albedo = sigs / sig;
	float final = 0.0f;
	Vector3 boxScale = Vector3(field.nx, field.ny, field.nz) * field.dx / (boxMax - boxMin);

	for(float t = t1; t >= t0; t -= dx)
	{
		Vector3 p = ray.origin() + ray.direction() * t;
		p = (p - boxMin) * boxScale;
		Array<Photon> photons;
		kdtree.getIntersectingMembers(Sphere(p, radius), photons);
		float rad = 0.0f, weight = 0.0;
		for(int i = 0; i < photons.size(); ++i)
		{
			// assume isotropic scattering
			float d2 = (photons[i].pos - p).squaredLength();
			float wp = exp(-d2 * r2r2);
			rad += photons[i].rad * wp;
			weight += wp;
		}
		rad *= albedo / (weight + 1e-6f);
		final = rad + final * exp(-sigdx * field.evaluate(p.x, p.y, p.z));
	}

	return final;
}

void VolumeRadianceTransfer::estimateRadianceSH(const Field<float> &field, const Vector3 &pos, SHFunction<float> &sh)
{
	Vector3 boxMin(0.0f, 0.0f, 0.0f);
	Vector3 boxMax(field.nx * field.dx, field.ny * field.dx, field.nz * field.dx);
	sh.values.assign(rbands * rbands, 0.0f);
	for(int l = 0; l < rbands; ++l) for(int m = -l; m <= l; ++m)
	{
		// Monte-Carlo integration over sphere
		float &coef = sh.values[l * l + l + m];
		for(int i = 0; i < rsamples; ++i)
		{
			float theta, phi;
			sh.random(theta, phi);
			float ylm = SHFunction<float>::realY(l, m, theta, phi);
			float cost = std::cos(theta);
			float sint = std::sin(theta);
			Ray ray(pos, Vector3(cos(phi) * sint, sin(phi) * sint, cost).direction());
			coef += estimateRadiance(field, ray, boxMin, boxMax) * ylm;
		}
		coef *= 4.0f * sh.PI / float(rsamples);
	}
}

void VolumeRadianceTransfer::renderScattering(const Field<float> &field, const Vector3 &boxMin, const Vector3 &boxMax, const GCamera &camera, Image1::Ref image)
{
	if(image->sizeInMemory() <= 0)
		return;

	Rect2D rect = Rect2D::xywh(0.0f, 0.0f, image->width(), image->height());
	for(int j = 0; j < image->height(); ++j) for(int i = 0; i < image->width(); ++i)
	{
		Ray ray = camera.worldRay(i + 0.5f, j + 0.5f, rect);
		float rad = estimateRadiance(field, ray, boxMin, boxMax);
		image->set(i, j, Color1(std::max(0.0f, rad)));
		//image.pixel4(i, j) = Color4(Color3::jetColorMap(clamp(rad * 0.5f + 0.5f, 0.0f, 1.0f)));
	}
}

float VolumeRadianceTransfer::sampleBoundaryRadiance(const Array<SHFunction<float>> &rad, int pitch, const Vector2 &pos, const Vector3 &dir)
{
	int n = pbands * pbands;
	int npitch = n * pitch;
	float theta = std::acos(dir.z);
	float cosphi = dir.x / sqrt(dir.x * dir.x + dir.y * dir.y);
	float phi = dir.y >= 0.0f ? std::acos(cosphi) : 2.0f * SHFunction<float>::PI - std::acos(cosphi);
	//float phi = std::acos(dir.x / std::sin(theta));
	int xi = floor(pos.x);
	int yi = floor(pos.y);
	float xf = pos.x - xi;
	float yf = pos.y - yi;
	float final = 0.0f;
	int idx = (xi * pitch + yi) * n;
	for(size_t i = 0; i < lighting.values.size(); ++i)
	{
		float r1 = rad[idx].evaluate(theta, phi);
		float r2 = rad[idx + n].evaluate(theta, phi);
		float r3 = rad[idx + npitch].evaluate(theta, phi);
		float r4 = rad[idx + npitch + n].evaluate(theta, phi);
		float r12 = r1 * (1.0f - yf) + r2 * yf;
		float r34 = r3 * (1.0f - yf) + r4 * yf;
		final += r12 * (1.0f - xf) + r34 * xf;
		++idx;
	}
	return final;
}

float VolumeRadianceTransfer::sampleBoundaryRadiance(const Field<float> &field, const Ray &ray, const Vector3 &boxMin, const Vector3 &boxMax)
{
	float t0, t1;
	bool hit = intersectRayBox(ray, boxMin, boxMax, t0, t1);
	if(!hit)
		return 0.0f;

	float br = 0.0f;
	Vector3 p = ray.origin() + ray.direction() * t0;
	p = (p - boxMin) / (boxMax - boxMin) * Vector3(field.nx, field.ny, field.nz);
	p = clamp(p, Vector3(0.501f, 0.501f, 0.501f), Vector3(field.nx - 0.501f, field.ny - 0.501f, field.nz - 0.501f)) - Vector3(0.5f, 0.5f, 0.5f);
	if(p.x < 0.5f)
		br = sampleBoundaryRadiance(bndrad[0], field.nz, p.yz(), ray.direction().direction());
	else if(p.x > field.nx - 1.5f)
		br = sampleBoundaryRadiance(bndrad[1], field.nz, p.yz(), ray.direction().direction());
	else if(p.y < 0.5f)
		br = sampleBoundaryRadiance(bndrad[2], field.nx, p.zx(), ray.direction().direction());
	else if(p.y > field.ny - 1.5f)
		br = sampleBoundaryRadiance(bndrad[3], field.nx, p.zx(), ray.direction().direction());
	else if(p.z < 0.5f)
		br = sampleBoundaryRadiance(bndrad[4], field.ny, p.xy(), ray.direction().direction());
	else if(p.z > field.nz - 1.5f)
		br = sampleBoundaryRadiance(bndrad[5], field.ny, p.xy(), ray.direction().direction());

	return br;
}

void VolumeRadianceTransfer::render(const Field<float> &field, const Vector3 &boxMin, const Vector3 &boxMax, const GCamera &camera, Image1::Ref image)
{
	if(image->sizeInMemory() <= 0)
		return;

	//image.convertToRGBA();
	Rect2D rect = Rect2D::xywh(0.0f, 0.0f, image->width(), image->height());
	for(int j = 0; j < image->height(); ++j) for(int i = 0; i < image->width(); ++i)
	{
		Ray ray = camera.worldRay(i + 0.5f, j + 0.5f, rect);
		float rad = sampleBoundaryRadiance(field, ray, boxMin, boxMax);
		image->set(i, j, Color1(std::max(0.0f, rad)));
		//image.pixel4(i, j) = Color4(Color3::jetColorMap(clamp(rad * 0.5f + 0.5f, 0.0f, 1.0f)));
	}
}


void VolumeRadianceTransfer::evaluateLightField(Field<float> &light)
{
	float radius = light.dx * 4.0f;
	//float radius = light.dx * 4;
	float r2r2 = 2.0f / (radius * radius);
	float rvol = 1.0f / (4.0f / 3.0f * SHFunction<float>::PI * 8.0f);
	for(int z = 0; z < light.nz; ++z) for(int y = 0; y < light.ny; ++y) for(int x = 0; x < light.nx; ++x)
	{
		Array<Photon> members;
		Vector3 pos = Vector3(x + 0.5f, y + 0.5f, z + 0.5f) * light.dx;
		kdtree.getIntersectingMembers(Sphere(pos, radius), members);
		light(x, y, z) = 0.0f;
		//float weight = 1e-6f;
		for(int i = 0; i < members.size(); ++i)
		{
			float d2 = (members[i].pos - pos).squaredLength();
			float wp = exp(-d2 * r2r2);
			light(x, y, z) += members[i].rad * wp;
			//weight += wp;
		}
		light(x, y, z) *= rvol;
		//light(x, y, z) /= weight;
	}
	light.syncHostToDevice();
}

void VolumeRadianceTransfer::renderVertexBuffer(RenderDevice *rd, const Vector3 &offset, const Vector3 &scale)
{
	if(kdtree.size() > 0)
	{
		//Draw::axes(CoordinateFrame(), rd);
		rd->pushState();
		glDisable(GL_LINE_SMOOTH);
		rd->setShadeMode(RenderDevice::SHADE_SMOOTH);
		rd->setColor(Color3::yellow());
		rd->beginPrimitive(RenderDevice::Primitive::LINES);
		for(PointKDTree<Photon>::Iterator it = kdtree.begin(); it != kdtree.end(); ++it)
		{
			rd->setColor(Color3::black());
			rd->sendVertex(offset + it->pos * scale);
			rd->setColor(Color3::jetColorMap((it->rad - minrad) / (maxrad - minrad)));
			rd->sendVertex(offset + (it->pos + it->dir * 0.5f) * scale);
		}
		rd->endPrimitive();
		glEnable(GL_LINE_SMOOTH);
		rd->popState();
		//pshader->args.set("MVP", rd->modelViewProjectionMatrix());
		//rd->pushState();
		//rd->setShader(pshader);
		//rd->setNormalArray(nb);
		//rd->beginIndexedPrimitives();
		//rd->setColor(Color3::white());
		//rd->setVertexArray(vb);
		//rd->sendSequentialIndices(RenderDevice::Primitive::POINTS, photons.size());
		//rd->endIndexedPrimitives();
		//rd->popState();
	}
}
